#!/bin/bash
#SBATCH --nodes=1
#SBATCH --partition=gpu
#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=96
#SBATCH --time=500:00:00
#SBATCH --job-name=equ_run

module load apps/gromacs/2020.6
module load apps/anaconda3/2021.11
source /public/software/apps/anaconda3/2021.11/bin/activate base
source /public/software/apps/anaconda3/2021.11/bin/activate amber

echo " Start equilibration"
function min_equ {
    [ ! -f $1/min01.tpr ] && { gmx grompp -c $1/system.pdb -r $1/system.pdb -f $1/min01.mdp -p $1/system.top -o $1/min01.tpr -n $1/index.ndx -po $1/min01_mdout.mdp -maxwarn 7 >$1/min01_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/min01.gro ] && { (cd $1/ && gmx mdrun --deffnm min01 -nt 1 >min01_mdrun_log.log 2>&1) || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/min02.tpr ] && { gmx grompp -c $1/min01.gro -r $1/min01.gro -f $1/min02.mdp -p $1/system.top -o $1/min02.tpr -n $1/index.ndx -po $1/min02_mdout.mdp -maxwarn 7 >$1/min02_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/min02.gro ] && { (cd $1/ && gmx mdrun --deffnm min02 -nt 1 >min02_mdrun_log.log 2>&1) || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/min03.tpr ] && { gmx grompp -c $1/min02.gro -r $1/min02.gro -f $1/min03.mdp -p $1/system.top -o $1/min03.tpr -n $1/index.ndx -po $1/min03_mdout.mdp -maxwarn 7 >$1/min03_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/min03.gro ] && { (cd $1/ && gmx mdrun --deffnm min03 -nt 1 >min03_mdrun_log.log 2>&1) || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/nve.tpr ] && { gmx grompp -c $1/min03.gro -r $1/min03.gro -f $1/nve.mdp -p $1/system.top -o $1/nve.tpr -n $1/index.ndx -po $1/nve_mdout.mdp -maxwarn 7 >$1/nve_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/nve.gro ] && { (cd $1/ && gmx mdrun --deffnm nve -nt 1 >nve_mdrun_log.log 2>&1) || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/nvt.tpr ] && { gmx grompp -c $1/nve.gro -r $1/nve.gro -f $1/nvt.mdp -p $1/system.top -o $1/nvt.tpr -n $1/index.ndx -po $1/nvt_mdout.mdp -maxwarn 7 >$1/nvt_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/nvt.gro ] && { (cd $1/ && gmx mdrun --deffnm nvt -nt 1 >nvt_mdrun_log.log 2>&1) || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/npt.tpr ] && { gmx grompp -c $1/nvt.gro -r $1/nvt.gro -f $1/npt.mdp -p $1/system.top -o $1/npt.tpr -n $1/index.ndx -po $1/npt_mdout.mdp -maxwarn 7 >$1/npt_grompp.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    [ ! -f $1/npt.gro ] && { (cd $1/ && gmx mdrun --deffnm npt -nt 1 >npt_mdrun_log.log 2>&1) || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
    # [ ! -f $1/md_nocont.tpr ] && { gmx grompp -c $1/npt.gro -r $1/npt.gro -t $1/npt.cpt -f $1/md_nocont.mdp -po $1/mdout_nocont.mdp -o $1/md_nocont.tpr -n $1/index.ndx -p $1/system.top -maxwarn 7 > $1/mdout_nocont.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; } }
    [ ! -f $1/md.tpr ] && { gmx grompp -c $1/npt.gro -r $1/npt.gro -t $1/npt.cpt -f $1/md.mdp -p $1/system.top -o $1/md.tpr -n $1/index.ndx -po $1/mdout.mdp -maxwarn 7 >$1/mdout.log 2>&1 || { echo "Failed to run command  at line ${LINENO}" && exit; }; }
}

n_jobs=0
max_jobs=90
pids=""

function waiting {
    for pid in $1; do
        wait $pid
    done
}

ligands=$(ls ligands/*.sdf)

for lig in $ligands; do
    lig_f_name=${lig:8}
    lig_name=${lig_f_name%%.*}
    echo ${lig_name}
    min_equ ${lig_name}_md &
    pids="$pids $!"
    n_jobs=$((n_jobs + 1))
    if (($n_jobs > max_jobs)); then
        waiting $pids
        n_jobs=0
        pids=""
    fi
done
waiting $pids

echo " Finished equilibration"

echo "Start run"
if [ $(ls *_md/md.gro 2>/dev/null | wc -l) -eq $(ls | egrep "_md+$" | wc -l) ]; then
    echo "[INFO] Run step already done. Going on."
else
    #[ $(ls *_md/lambda.gro 2>/dev/null | wc -l) -ne 0 ] && { echo "[ERROR] Some of the lambda windows are complete, but not all of them. Because this run uses HREX, I cannot continue"; exit 1; }
    mpirun -np $(ls | egrep "_md+$" | wc -l) mdrun_mpi -v --deffnm md -cpi md.cpt -multidir $(ls -v | egrep "_md+$") >&feprun_$(date "+%H%M%S_%d%m%Y").log || {
        echo "mdrun failed at line ${LINENO} "
        exit -1
    }
fi
